package org.sean.framework.web.filter;

import org.sean.framework.logging.Logger;
import org.sean.framework.util.DigestUtil;

import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletRequestWrapper;
import java.util.Collections;
import java.util.Enumeration;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Map;
import java.util.Set;
import java.util.UUID;

/**
 * 可变请求头
 */
public class MutableHeaderRequestWrapper extends HttpServletRequestWrapper {
    private final Map<String, String> headers;

    public MutableHeaderRequestWrapper(HttpServletRequest request) {
        super(request);
        this.headers = new HashMap<>();
    }

    public void putHeader(String name, String value) {
        this.headers.put(name, value);
    }

    @Override
    public String getHeader(String name) {
        if (Logger.KEY_TRACE.equals(name)) {
            String traceId = ((HttpServletRequest) getRequest()).getHeader(name);
            return traceId == null ? DigestUtil.md5Hex16(UUID.randomUUID().toString()) : traceId;
        }
        String headerValue = headers.get(name);

        if (headerValue != null) {
            return headerValue;
        }
        return ((HttpServletRequest) getRequest()).getHeader(name);
    }

    @Override
    public Enumeration<String> getHeaderNames() {
        if (headers.isEmpty()) {
            return super.getHeaderNames();
        }

        Set<String> set = new HashSet<>(headers.keySet());
        // 添加自定义header
        Enumeration<String> e = ((HttpServletRequest) getRequest()).getHeaderNames();
        while (e.hasMoreElements()) {
            String n = e.nextElement();
            set.add(n);
        }

        return Collections.enumeration(set);
    }
}
